"""
A collection of bands extraction EOTasks

Copyright (c) 2017- Sinergise and contributors
For the full list of contributors, see the CREDITS file in the root directory of this source tree.

This source code is licensed under the MIT license, see the LICENSE file in the root directory of this source tree.
"""

from __future__ import annotations

import numpy as np
import pytest

from sentinelhub import CRS, BBox

from eolearn.core import EOPatch, FeatureType
from eolearn.features import NormalizedDifferenceIndexTask

INPUT_FEATURE = (FeatureType.DATA, "TEST")


@pytest.mark.parametrize("bad_input", [(1, 2, 3), "test", 0.5])
def test_bad_input(bad_input):
    with pytest.raises(ValueError):
        NormalizedDifferenceIndexTask(INPUT_FEATURE, (FeatureType.DATA, "NDI"), bands=bad_input)


def test_ndi():
    eopatch = EOPatch(bbox=BBox((0, 0, 1, 1), CRS(3857)), timestamps=["1234-05-06"] * 4)
    eopatch[INPUT_FEATURE] = np.zeros((4, 3, 3, 9))

    band_a, band_b = 4.123, 3.321
    eopatch[INPUT_FEATURE][..., 0] = band_a
    eopatch[INPUT_FEATURE][..., 1] = band_b
    eopatch = NormalizedDifferenceIndexTask(INPUT_FEATURE, (FeatureType.DATA, "NDI"), bands=[0, 1]).execute(eopatch)
    assert (eopatch.data["NDI"] == ((band_a - band_b) / (band_a + band_b))).all()

    eopatch[INPUT_FEATURE][..., 5] = np.nan
    eopatch[INPUT_FEATURE][..., 7] = np.inf
    eopatch = NormalizedDifferenceIndexTask(INPUT_FEATURE, (FeatureType.DATA, "NAN_INF_INPUT"), bands=[5, 7]).execute(
        eopatch
    )
    assert np.isnan(eopatch.data["NAN_INF_INPUT"]).all()

    eopatch[INPUT_FEATURE][..., 1] = 1
    eopatch[INPUT_FEATURE][..., 3] = -1
    eopatch = NormalizedDifferenceIndexTask(INPUT_FEATURE, (FeatureType.DATA, "DIV_ZERO_NAN"), bands=[1, 3]).execute(
        eopatch
    )
    assert np.isnan(eopatch.data["DIV_ZERO_NAN"]).all()

    eopatch[INPUT_FEATURE][..., 1] = 0
    eopatch[INPUT_FEATURE][..., 3] = 0
    eopatch = NormalizedDifferenceIndexTask(
        INPUT_FEATURE, (FeatureType.DATA, "DIV_INVALID"), bands=[1, 3], undefined_value=123
    ).execute(eopatch)
    assert (eopatch.data["DIV_INVALID"] == 123).all()
